package com.feidee.fd.sml.algorithm.ml

import com.feidee.fd.sml.algorithm.component.ml.regression.{GBDTRegressionComponent, GBDTRegressionParam}
import com.feidee.fd.sml.algorithm.util.{TestingDataGenerator, ToolClass}
import org.scalatest.FunSuite

/**
  * @Author: dongguosheng
  * @Date: 2019/3/25 16:37
  * @Review songhaicheng
  * @Email: guosheng_dong@sui.com
  */
class GBDTRegressionComponentSuite extends FunSuite {

  val gbtr = new GBDTRegressionComponent()
  val paramStr: String =
    """
      |{
      |	'input_pt': '',
      |	'output_pt': '',
      |	'featuresCol': 'features',
      |	'labelCol': 'label',
      | 'probabilityCol': 'probability',
      |	'modelPath': '',
      |	'metrics': ['accuracy', 'auc', 'pr', 'abc'],
      | 'minInstancesPerNode': 1,
      |	'impurity': 'variance',
      |	'lossType': 'squared',
      |	'maxDepth': 5,
      |	'maxBins': 32,
      | 'minInfoGain': 0.0,
      |	'maxMemoryInMB': 256,
      |	'cacheNodeIds': false,
      |	'checkpointInterval': 10,
      |	'subsamplingRate': 1.0,
      | 'maxIter': 10,
      |	'stepSize': 0.1
      |}
    """.stripMargin

  val param:  GBDTRegressionParam = gbtr.parseParam(new ToolClass().encrypt(paramStr))

  /**
    * GBDT 参数构造方法测试：检测到读取出合法参数 & 成功赋值默认参数，即认为测试通过
    */
  test("GBDT parameter construction test") {
    assert("".equals(param.input_pt) & "".equals(param.output_pt) & "features".equals(param.featuresCol) &
      "label".equals(param.labelCol) & "squared".equals(param.lossType) &
      "variance".equals(param.impurity) & !param.cacheNodeIds & "".equals(param.modelPath) & param.maxDepth == 5 &
      param.metrics.sameElements(Array("accuracy", "auc", "pr", "abc")) & param.maxBins == 32 & param.subsamplingRate == 1.0
      & param.stepSize==0.1
    )
  }

  /**
    * GBDT 模型训练功能测试：传入生成的测试数据和正确参数，可以正常训练 GBDT 模型，且预测结果列等于生成数据列数，即认为测试通过
    */
  test("GBDT training function test") {
    // 测试模型训练方法
    val testingData = TestingDataGenerator.sampleIsotonicRegressionData
    testingData.show(20)
    val model = gbtr.train(param, testingData)
    //    assert(model.stages(2).asInstanceOf[GBTClassificationModel].numClasses == 2)
    //    assert(model.transform(testingData).count() == 20)
  }

}